package burp.Utils;

import burp.BurpExtender;
import burp.IHttpRequestResponse;
import burp.IResponseInfo;
import burp.ParamNullifierHistory;

import javax.swing.*;
import javax.swing.table.DefaultTableCellRenderer;
import java.awt.*;

public class StatusTableCellRenderer extends DefaultTableCellRenderer {


    @Override
    public Component getTableCellRendererComponent(JTable table, Object value, boolean isSelected, boolean hasFocus, int row, int column) {
        Component cellComponent = super.getTableCellRendererComponent(table, value, isSelected, hasFocus, row, column);

        // 获取 ParamNullifierHistory 实例
        ParamNullifierHistory history = (ParamNullifierHistory) table.getClientProperty("history");
        if (history == null) return cellComponent;

        // 获取当前行对应的请求对象
        IHttpRequestResponse message = history.getNullifiedRequest(row);
        if (message == null) return cellComponent;

        // 获取响应信息
        IResponseInfo resInfo = message.getResponse() != null ?
                BurpExtender.getInstance().getHelpers().analyzeResponse(message.getResponse()) : null;
        int statusCode = resInfo != null ? resInfo.getStatusCode() : 0;

        String responseStr = "";
        if (message.getResponse() != null) {
            responseStr = BurpExtender.getInstance().getHelpers().bytesToString(message.getResponse());
        }

        boolean shouldHighlight = false;

        // 判断是否应高亮整行
        if (BurpExtender.CHECK_RESPONSE_CODE && statusCode >= 500 && statusCode < 600) {
            shouldHighlight = true;
        } else {
            for (String keyword : BurpExtender.CUSTOM_RESPONSE_INDICATORS) {
                if (responseStr.contains(keyword)) {
                    shouldHighlight = true;
                    break;
                }
            }
        }

        // 如果条件满足，则整行标红
        if (shouldHighlight) {
            cellComponent.setBackground(Color.RED);
            cellComponent.setForeground(Color.WHITE);
        } else {
            cellComponent.setBackground(table.getBackground());
            cellComponent.setForeground(table.getForeground());
        }

        return cellComponent;
    }
}
